import os
os.environ['CUDA_VISIBLE_DEVICES']='0'

import torch,os,pdb,shutil
from diffusers import FluxControlPipeline,FluxPriorReduxPipeline
from diffusers.utils import load_image
from image_gen_aux import DepthPreprocessor

from util_flux import pad_image
from util_flux import horizontal_concat_images,process_img_1024

FLUX_REDUX='/home/shengjie/ckp/FLUX.1-Redux-dev'
FLUX_DEPTH='/home/shengjie/ckp/FLUX.1-Depth-dev'
FLUX_DEPTH_LORA='/home/shengjie/ckp/FLUX.1-Depth-dev-lora'
FLUX='/data/models/FLUX___1-dev'

DEPTH_PREDCITION='/home/shengjie/ckp/depth-anything-large-hf'

types = ['niukou','niukou-pockets','yinhua']
choose_type = types[0]
index = len(choose_type.split('-'))+1

# examples_dir = '/data/shengjie/style_zhenzhi/'
examples_dir = f'/mnt/nas/shengjie/datasets_zhenzhi/zhenzhi-{choose_type}/'
examples_dir2 = f'/mnt/nas/shengjie/datasets_zhenzhi/zhenzhi-{choose_type}-edited-restore/'
# save_dir = '/data/shengjie/synthesis_zhenzhi/'
save_dir = f'/mnt/nas/shengjie/datasets_zhenzhi/zhenzhi-{choose_type}-synthesis/'

# 临时处理
for name in  os.listdir(save_dir):
    if not name.endswith('.jpg'):continue
    try:
        t1,t2 = name.split('_')
        t2 = t2[:-4]

        if t1.split('-')[index] == t2.split('-')[index]:
            # pdb.set_trace()
            os.remove( os.path.join(save_dir,name) )
    except:
        pdb.set_trace()

print(f'还剩 {len(os.listdir(save_dir))}')
# pdb.set_trace()



if os.path.exists(save_dir):shutil.rmtree(save_dir)
os.makedirs(save_dir,)

target_shape = (1024,1024)

pipe_prior_redux = FluxPriorReduxPipeline.from_pretrained(
                                            FLUX_REDUX, 
                                            torch_dtype=torch.bfloat16).to("cuda")
pipe = FluxControlPipeline.from_pretrained(FLUX, torch_dtype=torch.bfloat16).to("cuda")
pipe.load_lora_weights(FLUX_DEPTH_LORA, adapter_name="depth")
pipe.set_adapters("depth", 0.7)  # 多大程度上的迁移？
processor = DepthPreprocessor.from_pretrained(DEPTH_PREDCITION)

'''
zhenzhi-{choose_type}                 原图
zhenzhi-{choose_type}-edited-restore  修复的扩展图

files1 = 原图 + 扩展图      全部
files2 = 原图               提供精准的style
'''
# ori + extend style
imagefiles_all = os.listdir(examples_dir) + os.listdir(examples_dir2)
# ori style
imagefiles = os.listdir(examples_dir)

# pdb.set_trace()

# test_img = os.path.join(examples_dir,imagefiles[0])
# test_img2 = os.path.join(examples_dir,imagefiles[1])
# for t1 in imagefiles:
#     for t2 in imagefiles:
from itertools import product
from tqdm import tqdm
# t1 遍历所有
# t2 只遍历 不同款式 '-' not in t2
# t1.prefix != t2.prefix

def synthesis_image(img_path1,img_path2):
    # img1 == depth
    # img2 == redux(embedding)

    # prompt = "A robot made of exotic candies and chocolates of different kinds. The background is filled with confetti and celebratory gifts."
    control_image = process_img_1024(img_path1)
    # control_image,_,_,_,_ = pad_image(control_image)
    # control_image = control_image.resize(target_shape)
    main_condition_image = process_img_1024(img_path2)
    # main_condition_image,_,_,_,_ = pad_image(main_condition_image)
    # main_condition_image = main_condition_image.resize(target_shape)

    # pdb.set_trace()

    control_image_depth = processor(control_image)[0].convert("RGB")

    prompt_emb,ppoled_prompt_emb = pipe_prior_redux(main_condition_image,
                                             return_dict=False) # attr 'prompt_embeds' torch.Size([1, 1241, 4096]) 
    # del pipe_prior_redux
    # torch.cuda.empty_cache()

    image = pipe(
        control_image=control_image_depth,
        height=control_image_depth.size[1],
        width=control_image_depth.size[0],
        num_inference_steps=20,
        guidance_scale=4.5,
        # generator=torch.Generator().manual_seed(42),
        prompt_embeds=prompt_emb,
        pooled_prompt_embeds=ppoled_prompt_emb,
    ).images[0]
    # image.save("output.png")
    # 横向拼接 control_image main_condition_image image
    horizontal_image = horizontal_concat_images([control_image,control_image_depth,
                            main_condition_image,image],)

    return image,horizontal_image    

count = 0
for t1,t2 in tqdm(product(imagefiles_all,imagefiles)):
    # xx-xx-xx.jpg split('-') 得到  [xx,xx,xx.jpg]
    # 所以要加 :-4
    if t1[:-4].split('-')[index] == t2[:-4].split('-')[index]:continue
   
    count += 1

    # pdb.set_trace()

    img_path1 = os.path.join(examples_dir,t1)
    if not os.path.exists(img_path1):
        img_path1 = os.path.join(examples_dir2,t1)
    img_path2 = os.path.join(examples_dir,t2)

    
    image,horizontal_image = synthesis_image(img_path1,img_path2)

    save_test_img = os.path.join(save_dir,
                                os.path.splitext(t1)[0]+'_'+\
                                    os.path.splitext(t2)[0] +\
                                    os.path.splitext(t1)[1]
                                )

    
    horizontal_image.save('tmp.jpg')

    # pdb.set_trace()
    image.save(save_test_img)
print('final 生成 : ', count )